﻿#region COPYRIGHT
//--------------------------------------------------------------------------------
// <copyright file="QuadTree.cs" company="starLiGHT Entertainment Studios">
//    Copyright (c) 2010, 2011 
//       Roland Rosenkranz (Glatzemann@email.de)
// </copyright>
// <license>
//   This file is part of starLiGHT.Core.
//
//   starLiGHT.Core is free software: you can redistribute it and/or modify
//   it under the terms of the GNU Lesser General Public License as published by
//   the Free Software Foundation, either version 3 of the License, or
//   (at your option) any later version.
//
//   starLiGHT.Core is distributed in the hope that it will be useful,
//   but WITHOUT ANY WARRANTY; without even the implied warranty of
//   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//   GNU Lesser General Public License for more details.
//
//   You should have received a copy of the GNU Lesser General Public License
//   along with starLiGHT.Core. If not, see http://www.gnu.org/licenses/.
//    
//   ADDITIONAL (commercial) LICENSES for starLiGHT.Core are available on request.
//
//
//
//   Original License:
//   http://csharpquadtree.codeplex.com/SourceControl/changeset/view/27798#506270
// 
//   modifications by Glatzemann so far:
//    - yield return
//    - switched to ReadOnlyCollection where possible (and needed)
//    - switched everywhere to float
//    - removed System.Windows reference constraint (used Vector2 instead of Size)
//    - ExpandRootNode: replace merging code with RectangleF Union code
// </license>
// <version>
// **************[ starLiGHT.Engine SVN ]**********************
// * $Rev::                       $:  Revision of last commit *
// * $Author::                    $:  Author of last commit   *
// * $Date::                      $:  Date of last commit     *
// ************************************************************
// </version>
//--------------------------------------------------------------------------------
#endregion

namespace starLiGHT.Collections
{
    #region Using Statements
    using System;
    using System.Collections;
    using System.Collections.Generic;
    using System.Collections.ObjectModel;
    using Microsoft.Xna.Framework;

    #endregion

    public enum Direction : int
    {
        NW = 0,
        NE = 1,
        SW = 2,
        SE = 3
    }

    public interface IQuadNode
    {
        event EventHandler BoundsChanged;

        RectangleF Bounds { get; }
    }

    public class QuadTree<T> : ICollection where T : class, IQuadNode
    {
        #region Member Variables
        private readonly bool sort;
        private readonly Vector2 minLeafSize;
        private readonly int maxObjectsPerLeaf;
        private QuadNode root = null;
        private Dictionary<T, QuadNode> objectToNodeLookup = new Dictionary<T, QuadNode>();
        private Dictionary<T, int> objectSortOrder = new Dictionary<T, int>();
        private object syncLock = new object();
        private int objectSortId = 0;
        private FastList<T> results = new FastList<T>();

        #endregion

        #region Constructors
        /// <summary>
        /// Initializes a new instance of the QuadTree class.
        /// </summary>
        /// <param name="minLeafSize">The smallest size a leaf will split into</param>
        /// <param name="maxObjectsPerLeaf">Maximum number of objects per leaf before it forces a split into sub quadrants</param>
        public QuadTree(Vector2 minLeafSize, int maxObjectsPerLeaf)
        {
            this.minLeafSize = minLeafSize;
            this.maxObjectsPerLeaf = maxObjectsPerLeaf;
        }

        /// <summary>
        /// Initializes a new instance of the QuadTree class.
        /// </summary>
        /// <param name="minLeafSize">The smallest size a leaf will split into</param>
        /// <param name="maxObjectsPerLeaf">Maximum number of objects per leaf before it forces a split into sub quadrants</param>
        /// <param name="sort">Whether or not queries will return objects in the order in which they were added</param>
        public QuadTree(Vector2 minLeafSize, int maxObjectsPerLeaf, bool sort)
            : this(minLeafSize, maxObjectsPerLeaf)
        {
            this.sort = sort;
        }

        #endregion

        #region Properties
        public int Count
        {
            get
            {
                lock (this.syncLock)
                {
                    if (this.root == null)
                    {
                        return 0;
                    }

                    return this.GetQuadObjectCount(this.root);
                }
            }
        }

        public bool IsSynchronized
        {
            get
            {
                return true;
            }
        }

        public object SyncRoot
        {
            get
            {
                throw new NotImplementedException();
            }
        }

        public QuadNode Root
        {
            get 
            { 
                return this.root; 
            }
        }

        public Vector2 MinLeafSize
        {
            get 
            { 
                return this.minLeafSize; 
            }
        }

        public int MaxObjectsPerLeaf
        {
            get 
            { 
                return this.maxObjectsPerLeaf; 
            }
        }

        public bool Sort
        {
            get 
            { 
                return this.sort; 
            }
        }

        #endregion

        #region Methods
        public int GetSortOrder(T quadObject)
        {
            lock (this.objectSortOrder)
            {
                if (!this.objectSortOrder.ContainsKey(quadObject))
                {
                    return -1;
                }
                else
                {
                    return this.objectSortOrder[quadObject];
                }
            }
        }

        public void Insert(T quadObject)
        {
            lock (this.syncLock)
            {
                if (this.sort & !this.objectSortOrder.ContainsKey(quadObject))
                {
                    this.objectSortOrder.Add(quadObject, this.objectSortId++);
                }

                RectangleF bounds = quadObject.Bounds;
                if (this.root == null)
                {
                    Vector2 rootSize = new Vector2((float)Math.Ceiling(bounds.Width / this.minLeafSize.X), (float)Math.Ceiling(bounds.Height / this.minLeafSize.Y));
                    float multiplier = Math.Max(rootSize.X, rootSize.Y);
                    rootSize = new Vector2(this.minLeafSize.X * multiplier, this.minLeafSize.Y * multiplier);
                    Vector2 center = new Vector2(bounds.X + (bounds.Width * 0.5f), bounds.Y + (bounds.Height * 0.5f));
                    Vector2 rootOrigin = new Vector2(center.X - (rootSize.X * 0.5f), center.Y - (rootSize.Y * 0.5f));
                    this.root = new QuadNode(new RectangleF(rootOrigin, rootSize));
                }

                while (!this.root.Bounds.Contains(bounds))
                {
                    this.ExpandRoot(bounds);
                }

                this.InsertNodeObject(this.root, quadObject);
            }
        }

        public ReadOnlyCollection<T> Query(RectangleF bounds)
        {
            lock (this.syncLock)
            {
                this.results.Clear();
                if (this.root != null)
                {
                    this.Query(bounds, this.root, this.results);
                }

                if (this.sort)
                {
                    // this.results.Sort((a, b) => { return objectSortOrder[a].CompareTo(objectSortOrder[b]); });
                    throw new NotImplementedException();
                }

                return new ReadOnlyCollection<T>(this.results);
            }
        }

        public void Remove(T quadObject)
        {
            lock (this.syncLock)
            {
                if (this.sort && this.objectSortOrder.ContainsKey(quadObject))
                {
                    this.objectSortOrder.Remove(quadObject);
                }

                if (!this.objectToNodeLookup.ContainsKey(quadObject))
                {
                    throw new KeyNotFoundException("QuadObject not found in dictionary for removal");
                }

                QuadNode containingNode = this.objectToNodeLookup[quadObject];
                this.RemoveQuadObjectFromNode(quadObject);

                if (containingNode.Parent != null)
                {
                    this.CheckChildNodes(containingNode.Parent);
                }
            }
        }

        public int GetQuadNodeCount()
        {
            lock (this.syncLock)
            {
                if (this.root == null)
                {
                    return 0;
                }

                int count = this.GetQuadNodeCount(this.root, 1);

                return count;
            }
        }

        public IEnumerable<QuadNode> GetAllNodes()
        {
            lock (this.syncLock)
            {
                if (this.root != null)
                {
                    yield return this.root;

                    foreach (QuadNode childNode in this.GetChildNodes(this.root))
                    {
                        yield return childNode;
                    }
                }
            }
        }

        public void CopyTo(Array array, int index)
        {
            throw new NotImplementedException();
        }

        public IEnumerator GetEnumerator()
        {
            throw new NotImplementedException();
        }

        #endregion

        #region Private Members

        private void Query(RectangleF bounds, QuadNode node, FastList<T> results)
        {
            lock (this.syncLock)
            {
                if (node == null)
                {
                    return;
                }

                if (bounds.Intersects(node.Bounds))
                {
                    foreach (T quadObject in node.Objects)
                    {
                        if (bounds.Intersects(quadObject.Bounds))
                        {
                            results.Add(quadObject);
                        }
                    }

                    foreach (QuadNode childNode in node.Nodes)
                    {
                        this.Query(bounds, childNode, results);
                    }
                }
            }
        }

        private void ExpandRoot(RectangleF newChildBounds)
        {
            lock (this.syncLock)
            {
                bool isNorth = this.root.Bounds.Y < newChildBounds.Y;
                bool isWest = this.root.Bounds.X < newChildBounds.X;

                Direction rootDirection;
                if (isNorth)
                {
                    rootDirection = isWest ? Direction.NW : Direction.NE;
                }
                else
                {
                    rootDirection = isWest ? Direction.SW : Direction.SE;
                }

                RectangleF newRootBounds = RectangleF.Union(this.root.Bounds, newChildBounds);

                QuadNode newRoot = new QuadNode(newRootBounds);
                this.SetupChildNodes(newRoot);
                newRoot[rootDirection] = this.root;
                this.root = newRoot;
            }
        }

        private void InsertNodeObject(QuadNode node, T quadObject)
        {
            lock (this.syncLock)
            {
                if (!node.Bounds.Contains(quadObject.Bounds))
                {
                    throw new Exception("This should not happen, child does not fit within node bounds");
                }

                if (!node.HasChildNodes() && node.Objects.Count + 1 > this.maxObjectsPerLeaf)
                {
                    this.SetupChildNodes(node);

                    List<T> childObjects = new List<T>(node.Objects);
                    List<T> childrenToRelocate = new List<T>();

                    foreach (T childObject in childObjects)
                    {
                        foreach (QuadNode childNode in node.Nodes)
                        {
                            if (childNode == null)
                            {
                                continue;
                            }

                            if (childNode.Bounds.Contains(childObject.Bounds))
                            {
                                childrenToRelocate.Add(childObject);
                            }
                        }
                    }

                    foreach (T childObject in childrenToRelocate)
                    {
                        this.RemoveQuadObjectFromNode(childObject);
                        this.InsertNodeObject(node, childObject);
                    }
                }

                foreach (QuadNode childNode in node.Nodes)
                {
                    if (childNode != null)
                    {
                        if (childNode.Bounds.Contains(quadObject.Bounds))
                        {
                            this.InsertNodeObject(childNode, quadObject);
                            return;
                        }
                    }
                }

                this.AddQuadObjectToNode(node, quadObject);
            }
        }

        private void ClearQuadObjectsFromNode(QuadNode node)
        {
            lock (this.syncLock)
            {
                List<T> quadObjects = new List<T>(node.Objects);
                foreach (T quadObject in quadObjects)
                {
                    this.RemoveQuadObjectFromNode(quadObject);
                }
            }
        }

        private void RemoveQuadObjectFromNode(T quadObject)
        {
            lock (this.syncLock)
            {
                QuadNode node = this.objectToNodeLookup[quadObject];
                node.QuadObjects.Remove(quadObject);
                this.objectToNodeLookup.Remove(quadObject);
                quadObject.BoundsChanged -= new EventHandler(this.quadObject_BoundsChanged);
            }
        }

        private void AddQuadObjectToNode(QuadNode node, T quadObject)
        {
            lock (this.syncLock)
            {
                node.QuadObjects.Add(quadObject);
                this.objectToNodeLookup.Add(quadObject, node);
                quadObject.BoundsChanged += new EventHandler(this.quadObject_BoundsChanged);
            }
        }

        private void quadObject_BoundsChanged(object sender, EventArgs e)
        {
            lock (this.syncLock)
            {
                T quadObject = sender as T;
                if (quadObject != null)
                {
                    QuadNode node = this.objectToNodeLookup[quadObject];
                    if (!node.Bounds.Contains(quadObject.Bounds) || node.HasChildNodes())
                    {
                        this.RemoveQuadObjectFromNode(quadObject);
                        this.Insert(quadObject);
                        if (node.Parent != null)
                        {
                            this.CheckChildNodes(node.Parent);
                        }
                    }
                }
                else
                {
                    throw new Exception("The sender of BoundsChanged have to be of type " + typeof(T).ToString());
                }
            }
        }

        private void SetupChildNodes(QuadNode node)
        {
            lock (this.syncLock)
            {
                if (this.minLeafSize.X <= node.Bounds.Width / 2 && this.minLeafSize.Y <= node.Bounds.Height / 2)
                {
                    node[Direction.NW] = new QuadNode(node.Bounds.X, node.Bounds.Y, node.Bounds.Width / 2, node.Bounds.Height / 2);
                    node[Direction.NE] = new QuadNode(node.Bounds.X + (node.Bounds.Width / 2), node.Bounds.Y, node.Bounds.Width / 2, node.Bounds.Height / 2);
                    node[Direction.SW] = new QuadNode(node.Bounds.X, node.Bounds.Y + (node.Bounds.Height / 2), node.Bounds.Width / 2, node.Bounds.Height / 2);
                    node[Direction.SE] = new QuadNode(node.Bounds.X + (node.Bounds.Width / 2), node.Bounds.Y + (node.Bounds.Height / 2), node.Bounds.Width / 2, node.Bounds.Height / 2);
                }
            }
        }

        private void CheckChildNodes(QuadNode node)
        {
            lock (this.syncLock)
            {
                if (this.GetQuadObjectCount(node) <= this.maxObjectsPerLeaf)
                {
                    // Move child objects into this node, and delete sub nodes
                    List<T> subChildObjects = this.GetChildObjects(node);
                    foreach (T childObject in subChildObjects)
                    {
                        if (!node.Objects.Contains(childObject))
                        {
                            this.RemoveQuadObjectFromNode(childObject);
                            this.AddQuadObjectToNode(node, childObject);
                        }
                    }

                    if (node[Direction.NW] != null)
                    {
                        node[Direction.NW].Parent = null;
                        node[Direction.NW] = null;
                    }

                    if (node[Direction.NE] != null)
                    {
                        node[Direction.NE].Parent = null;
                        node[Direction.NE] = null;
                    }

                    if (node[Direction.SW] != null)
                    {
                        node[Direction.SW].Parent = null;
                        node[Direction.SW] = null;
                    }

                    if (node[Direction.SE] != null)
                    {
                        node[Direction.SE].Parent = null;
                        node[Direction.SE] = null;
                    }

                    if (node.Parent != null)
                    {
                        this.CheckChildNodes(node.Parent);
                    }
                    else
                    {
                        // Its the root node, see if we're down to one quadrant, with none in local storage - if so, ditch the other three
                        int numQuadrantsWithObjects = 0;
                        QuadNode nodeWithObjects = null;
                        foreach (QuadNode childNode in node.Nodes)
                        {
                            if (childNode != null && this.GetQuadObjectCount(childNode) > 0)
                            {
                                numQuadrantsWithObjects++;
                                nodeWithObjects = childNode;
                                if (numQuadrantsWithObjects > 1)
                                {
                                    break;
                                }
                            }
                        }

                        if (numQuadrantsWithObjects == 1)
                        {
                            foreach (QuadNode childNode in node.Nodes)
                            {
                                if (childNode != nodeWithObjects)
                                {
                                    childNode.Parent = null;
                                }
                            }

                            this.root = nodeWithObjects;
                        }
                    }
                }
            }
        }

        private List<T> GetChildObjects(QuadNode node)
        {
            lock (this.syncLock)
            {
                List<T> results = new List<T>();
                results.AddRange(node.QuadObjects);
                foreach (QuadNode childNode in node.Nodes)
                {
                    if (childNode != null)
                    {
                        results.AddRange(this.GetChildObjects(childNode));
                    }
                }

                return results;
            }
        }

        private int GetQuadObjectCount(QuadNode node)
        {
            lock (this.syncLock)
            {
                int count = node.Objects.Count;
                foreach (QuadNode childNode in node.Nodes)
                {
                    if (childNode != null)
                    {
                        count += this.GetQuadObjectCount(childNode);
                    }
                }

                return count;
            }
        }

        private int GetQuadNodeCount(QuadNode node, int count)
        {
            lock (this.syncLock)
            {
                if (node == null)
                {
                    return count;
                }

                foreach (QuadNode childNode in node.Nodes)
                {
                    if (childNode != null)
                    {
                        count++;
                    }
                }

                return count;
            }
        }

        private IEnumerable<QuadNode> GetChildNodes(QuadNode node)
        {
            // no lock needed -> it's private and called only from already locked
            // parent method.
            foreach (QuadNode childNode in node.Nodes)
            {
                if (childNode != null)
                {
                    if (childNode.HasChildNodes())
                    {
                        foreach (QuadNode childChildNode in this.GetChildNodes(childNode))
                        {
                            yield return childChildNode;
                        }
                    }
                    else
                    {
                        yield return childNode;
                    }
                }
            }
        }

        #endregion

        #region Class QuadNode
        public class QuadNode
        {
            public readonly int ID = id++;

            private static int id = 0;
            private QuadNode[] nodes = new QuadNode[4];
            private ReadOnlyCollection<T> objects;
            private ReadOnlyCollection<QuadNode> readOnlyNodes;
            private List<T> quadObjects = new List<T>();

            public QuadNode(RectangleF bounds)
            {
                this.Bounds = bounds;
                this.readOnlyNodes = new ReadOnlyCollection<QuadNode>(this.nodes);
                this.objects = new ReadOnlyCollection<T>(this.quadObjects);
            }

            public QuadNode(float x, float y, float width, float height)
                : this(new RectangleF(x, y, width, height))
            {
            }

            public List<T> QuadObjects
            {
                get
                {
                    return this.quadObjects;
                }
            }

            public ReadOnlyCollection<QuadNode> Nodes
            {
                get
                {
                    return this.readOnlyNodes;
                }
            }

            public ReadOnlyCollection<T> Objects
            {
                get
                {
                    return this.objects;
                }
            }

            public QuadNode Parent 
            { 
                get; 
                internal set; 
            }

            public RectangleF Bounds
            {
                get;
                internal set;
            }

            public QuadNode this[Direction direction]
            {
                get
                {
                    switch (direction)
                    {
                        case Direction.NW:
                            return this.nodes[0];
                        case Direction.NE:
                            return this.nodes[1];
                        case Direction.SW:
                            return this.nodes[2];
                        case Direction.SE:
                            return this.nodes[3];
                        default:
                            return null;
                    }
                }

                set
                {
                    switch (direction)
                    {
                        case Direction.NW:
                            this.nodes[0] = value;
                            break;
                        case Direction.NE:
                            this.nodes[1] = value;
                            break;
                        case Direction.SW:
                            this.nodes[2] = value;
                            break;
                        case Direction.SE:
                            this.nodes[3] = value;
                            break;
                    }

                    if (value != null)
                    {
                        value.Parent = this;
                    }
                }
            }

            public bool HasChildNodes()
            {
                return this.nodes[0] != null;
            }
        }

        #endregion
    }
}